#!/bin/python3

import codecs
import re
import xml.etree.ElementTree as xml
import sys

def prettyXml(element, indent, newline, level = 0): 
  if element:    
    if element.text == None or element.text.isspace():     
      element.text = newline + indent * (level + 1)      
    else:    
      element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)    
  temp = list(element)
  for subelement in temp:    
    if temp.index(subelement) < (len(temp) - 1):     
      subelement.tail = newline + indent * (level + 1)    
    else: 
      subelement.tail = newline + indent * level   
    prettyXml(subelement, indent, newline, level = level + 1)

def byteToText(bt):
  i = 0
  ret = ""
  while True:
    temp = bt[i:i+2]
    if temp == b'\xFF\xFF':
      break

    temp_str = ""
    try:
      temp_str = codecs.decode(temp, 'sjis')
      assert(len(temp_str) == 1)

      # 全角转半角
      uchar = ord(temp_str)
      if (uchar >= 0xFF01 and uchar <= 0xFF5E):
        uchar = uchar - 0xFEE0
      elif uchar == 0x3000:
        uchar = 0x20
      temp_str = chr(uchar)
    except:
      temp_str = '{' + temp.hex() + '}'

    ret = ret + temp_str
    i = i + 2
  return ret

def textToByte(txt):
  ret = b''
  i = 0
  while i < len(txt):
    c = txt[i]
    i = i + 1

    if c == '{':
      temp = txt[i-1:i+5]
      if re.match('\{[0-9a-f]{4}\}', temp):
        ret = ret + bytes.fromhex(temp[1:5])
        i = i + 5
        continue

    # 半角转全角
    uchar = ord(c)
    if (uchar >= 0x21 and uchar <= 0x7E):
      uchar = uchar + 0xFEE0
    elif uchar == 0x20:
      uchar = 0x3000
    c = chr(uchar)

    ret = ret + codecs.encode(c, 'sjis')
  ret = ret + b'\xFF\xFF'
  return ret

class MBM:
  def __init__(self, filename):
    self.filename = filename
    self.entries = []
    fp = open(filename, 'rb')
    self.readHeader(fp)
    self.readEntries(fp)
    self.readText(fp)
    fp.close()

  def readHeader(self, fp):
    fp.read(4)
    self.magic = fp.read(4)
    self.const0 = fp.read(4)
    self.fsize = int.from_bytes(fp.read(4), 'little')
    self.entriesCount = int.from_bytes(fp.read(4), 'little')
    self.entriesOffset = int.from_bytes(fp.read(4), 'little')
    fp.read(8)

  def readEntries(self, fp):
    i = 0
    while i < self.entriesCount:
      ent = self.readEntry(fp)
      if ent["offset"] != 0:
        i = i + 1

  def readEntry(self, fp):
    ent = {}
    ent["id"] = int.from_bytes(fp.read(4), 'little')
    ent["size"] = int.from_bytes(fp.read(4), 'little')
    ent["offset"] = int.from_bytes(fp.read(4), 'little')
    ent["padding"] = fp.read(4)
    self.entries.append(ent)
    return ent

  def readText(self, fp):
    for ent in self.entries:
      if ent["offset"] == 0:
        continue
      fp.seek(ent["offset"])
      txt = fp.read(ent["size"])
      ent["text"] = byteToText(txt)

  def saveAsXML(self):
    savename = self.filename[0:len(self.filename) - 3] + 'xml'
    root = xml.Element('mbm')
    for ent in self.entries:
      if ent["offset"] == 0:
        continue
      entry = xml.SubElement(root, 'entry')
      entry.set("id", str(ent["id"]))
      source = xml.SubElement(entry, 'source')
      source.text = ent['text']
      target = xml.SubElement(entry, 'target')
      target.text = ent['text']
    prettyXml(root, '  ', '\n')
    tree = xml.ElementTree(root)
    tree.write(savename, encoding='utf-8')

  def loadXML(self, xmlfile):
    src = xml.ElementTree(file = xmlfile)
    src_items = src.getroot().findall("./")
    fsize = 0
    for item in src_items:
      id = int(item.get('id'))
      ent = self.entries[int(id)]
      target = item.find('./target').text
      bt = textToByte(target)
      ent['target'] = bt
      ent['size'] = len(bt)
      fsize = fsize + len(bt)
      if id > 0:
        while True:
          id = id - 1
          e = self.entries[id]
          if e['offset'] == 0:
            continue
          ent['offset'] = e['offset'] + e['size']
          break
    fsize = fsize + self.entriesCount * 16 + self.entriesOffset
    self.fsize = fsize
    
  def saveAsMBM(self):
    savename = self.filename
    fp = open(savename, 'wb')
    self.writeHeader(fp)
    self.writeEntries(fp)
    self.writeText(fp)
    fp.close()

  def writeHeader(self, fp):
    fp.write(int.to_bytes(0, 4, 'little'))
    fp.write(self.magic)
    fp.write(self.const0)
    fp.write(int.to_bytes(self.fsize, 4, 'little'))
    fp.write(int.to_bytes(self.entriesCount, 4, 'little'))
    fp.write(int.to_bytes(self.entriesOffset, 4, 'little'))
    fp.write(int.to_bytes(0, 8, 'little'))

  def writeEntries(self, fp):
    for ent in self.entries:
      fp.write(int.to_bytes(ent["id"], 4, 'little'))
      fp.write(int.to_bytes(ent["size"], 4, 'little'))
      fp.write(int.to_bytes(ent["offset"], 4, 'little'))
      fp.write(ent["padding"])

  def writeText(self, fp):
    for ent in self.entries:
      try:
        ent['target']
      except:
        continue
      fp.write(ent['target'])


def main():
  args = sys.argv
  if len(args) == 1:
    print("Usage:")
    print(sys.argv[0], "export <mbm file>")
    print(sys.argv[0], "import <mbm file> <xml file>")
    return

  if args[1] == "export":
    mbm = MBM(args[2])
    mbm.saveAsXML()
  elif args[1] == "import":
    mbm = MBM(args[2])
    mbm.loadXML(args[3])
    mbm.saveAsMBM()
  else:
    print("Unknown command", args[1])

if __name__ == "__main__":
  main()
